
#include "Form.h"
#include "CalcT.h"
#include "BinData2.h"

template <class T> inline T SQR(T x) { return x*x; }

void WriteEE(std::ostream& fout, double minsep, double binsize, double smoothscale,
             const std::vector<BinData2<EData,EData> >& data)
{
    Form sci;
    Form fix; 
    sci.sci().prec(3).width(11).right().trail(1); 
    fix.fix().prec(4).width(11).right().trail(1);

    fout <<
        "# R_nominal .    <R>    "
        ".    xi+    .    xi-    .   xi+_im  .   xi-_im  .   sig_xi  "
        ".   weight  .   npairs  ";
    if (smoothscale > 0.)
        fout << ".    R_sm   .   xi+_sm  .   xi-_sm  . sig_xi_sm ";
    fout << std::endl;

    for(int i=0;i<int(data.size());++i) {

        double R = exp(data[i].meanlogr);
        fout << fix(minsep*exp((i+0.5)*binsize)) << fix(R);
        fout << sci(real(data[i].xiplus));
        fout << sci(real(data[i].ximinus));
        fout << sci(imag(data[i].xiplus));
        fout << sci(imag(data[i].ximinus));
        Assert(data[i].varxi >= 0.);
        fout << sci(sqrt(data[i].varxi));
        fout << sci(data[i].weight);
        fout << sci(data[i].npair);

        if (smoothscale > 0.) {
            std::complex<double> xiplussm=0.,ximinussm=0.;
            double varxism=0.,weightsm=0.,meanlogrsm=0.;

            for(int j=0;j<int(data.size());++j) {

                double r = exp(data[j].meanlogr);
                double xip = real(data[j].xiplus);
                double xim = real(data[j].ximinus);
                double var = data[j].varxi;
                double wj = data[j].weight;
                double s = r/R;

                if (s>1/smoothscale && s<smoothscale) {
                    meanlogrsm += data[j].meanlogr*wj;
                    xiplussm += xip*wj;
                    ximinussm += xim*wj;
                    varxism += var*SQR(wj);
                    weightsm += wj;
                }
            }

            if (weightsm > 0) {
                meanlogrsm /= weightsm;
                xiplussm /= weightsm;
                ximinussm /= weightsm;
                varxism /= SQR(weightsm);
            }

            fout << fix((weightsm==0.) ? R : exp(meanlogrsm));
            fout << sci(real(xiplussm));
            fout << sci(real(ximinussm));
            Assert(varxism >= 0.);
            fout << sci(sqrt(varxism));
        }
        fout << std::endl;
    }
}

void WriteM2(std::ostream& fout, double minsep, double binsize,
             const std::vector<BinData2<EData,EData> >& data)
{
    Form sci;
    Form fix; 
    sci.sci().prec(3).width(11).right().trail(1); 
    fix.fix().prec(4).width(11).right().trail(1);

    // Use s = r/R:
    // <Map^2>(R) = int_r=0..2R [s^2 dlogr (T+(s) xi+(r) + T-(s) xi-(r))/2]
    // <Mx^2>(R)  =     "       [ "    "   (T+(s) xi+(r) - T-(s) xi-(r))/2]
    // <Gam^2>(R) = int_r=0..2R [s^2 dlogr S+(s) xi+(r)]
    // <Gam^2>(E/B)(R) = int_r=0..2R [s^2 dlogr (S+(s) xi+(r) +- S+(s) xi-(r))/2]

    const double dlogr = binsize;

    fout <<
        "# R_nominal .    <R>    "
        ".  <Map^2>  .  sig_map  .   <Mx^2>  .   sig_mx  .   <MMx>   .  sig_mmx  "
        ".   <zero>  . sig_zero  .  <Gam^2>  .  sig_gam\n";
    for(int i=0;i<int(data.size());++i) {
        double mapsq=0.,mxsq=0.,varmap=0.,mmx=0.,zero=0.;
        double gamsq=0.,vargam=0.;
        //double gamEsq=0.,gamBsq=0.,varEBgam=0.;

        double R = exp(data[i].meanlogr);

        for(int j=0;j<int(data.size());++j) {

            double r = exp(data[j].meanlogr);
            double xip = real(data[j].xiplus);
            double xim = real(data[j].ximinus);
            double xipi = imag(data[j].xiplus);
            double ximi = imag(data[j].ximinus);
            double var = data[j].varxi;
            double s = r/R;
            double ssqdlogr = s*s*dlogr;
            double tp = Tplus(s), tm = Tminus(s);
            double sp = Splus(s);
            //double sm = Sminus(s);

            mapsq += ssqdlogr*(tp*xip + tm*xim)/2.;
            mxsq  += ssqdlogr*(tp*xip - tm*xim)/2.;
            mmx   += ssqdlogr*(tp*xipi + tm*ximi)/2.;
            zero  += ssqdlogr*(tp*xipi - tm*ximi)/2.;
            varmap += SQR(ssqdlogr)*(SQR(tp)+SQR(tm))*var/4.;

            gamsq += ssqdlogr * sp * xip;
            vargam += SQR(ssqdlogr*sp)*var;
            //gamEsq += ssqdlogr * (sp*xip + sm*xim)/2.;
            //gamBsq += ssqdlogr * (sp*xip - sm*xim)/2.;
            //varEBgam += SQR(ssqdlogr)*(SQR(sp)+SQR(sm))*var/4.;
        }

        fout << fix(minsep*exp((i+0.5)*binsize)) << fix(R);
        Assert(varmap >= 0.);
        fout << sci(mapsq) << sci(sqrt(varmap));
        fout << sci(mxsq) << sci(sqrt(varmap));
        fout << sci(mmx) << sci(sqrt(varmap));
        fout << sci(zero) << sci(sqrt(varmap));
        Assert(vargam >= 0.);
        fout << sci(gamsq) << sci(sqrt(vargam));
        fout << std::endl;
    }
}

void WriteNE(std::ostream& fout, double minsep, double binsize, double smoothscale,
             const std::vector<BinData2<NData,EData> >& data)
{
    Form sci;
    Form fix; 
    sci.sci().prec(3).width(11).right().trail(1); 
    fix.fix().prec(4).width(11).right().trail(1);

    fout <<
        "# R_nominal .    <R>    "
        ".  <gamT>   .   <gamX>  .   sig     "
        ".   weight  .   npairs  ";
    if (smoothscale > 0.)
        fout << ".    R_sm   .   gamT_sm . sig_sm    ";
    fout << std::endl;

    for(int i=0;i<int(data.size());++i) {

        double R = exp(data[i].meanlogr);
        fout << fix(minsep*exp((i+0.5)*binsize)) << fix(R);
        fout << sci(real(data[i].meangammat));
        fout << sci(imag(data[i].meangammat));
        Assert(data[i].vargammat >= 0.);
        fout << sci(sqrt(data[i].vargammat));
        fout << sci(data[i].weight);
        fout << sci(data[i].npair);

        if (smoothscale > 0.) {
            std::complex<double> gamsm=0.;
            double vargamsm=0.,weightsm=0.,meanlogrsm=0.;

            for(int j=0;j<int(data.size());++j) {

                double r = exp(data[j].meanlogr);
                double gam = real(data[j].meangammat);
                double var = data[j].vargammat;
                double wj = data[j].weight;
                double s = r/R;

                if (s>1/smoothscale && s<smoothscale) {
                    meanlogrsm += data[j].meanlogr*wj;
                    gamsm += gam*wj;
                    vargamsm += var*SQR(wj);
                    weightsm += wj;
                }
            }

            if (weightsm > 0) {
                meanlogrsm /= weightsm;
                gamsm /= weightsm;
                vargamsm /= SQR(weightsm);
            }

            fout << fix((weightsm==0.) ? R : exp(meanlogrsm));
            fout << sci(real(gamsm));
            Assert(vargamsm >= 0.);
            fout << sci(sqrt(vargamsm));
        }
        fout << std::endl;
    }
}

void WriteNM(std::ostream& fout, double minsep, double binsize, double smoothscale,
             const std::vector<BinData2<NData,EData> >& crossdata)
{
    Form sci;
    Form fix; 
    sci.sci().prec(3).width(11).right().trail(1); 
    fix.fix().prec(4).width(11).right().trail(1);

    double dlogr = binsize;

    fout <<
        "# R_nominal .    <R>    "
        ".  <NMap>   .  sig_nmap .   <NMx>   .  sig_nmx  "
        ".  <gamT>   .   sig_gT  .   <gamX>  .   sig_gX \n";
    for(int i=0;i<int(crossdata.size());++i) {

        double R = exp(crossdata[i].meanlogr);
        double nmap=0., nmx=0., varnm=0.;

        for(int j=0;j<int(crossdata.size());++j) {

            double r = exp(crossdata[j].meanlogr);
            double s = r/R;
            double ssqdlogr = s*s*dlogr;
            double gt = real(crossdata[j].meangammat);
            double gx = imag(crossdata[j].meangammat);
            double tc = Tcross(s);
            nmap += ssqdlogr * tc * gt;
            nmx += ssqdlogr * tc * gx;
            varnm += SQR(ssqdlogr*tc)*crossdata[j].vargammat;
        }
        fout << fix(minsep*exp((i+0.5)*binsize)) << fix(R);
        Assert(varnm >= 0.);
        fout << sci(nmap) << sci(sqrt(varnm));
        fout << sci(nmx) << sci(sqrt(varnm));

        if (smoothscale > 0.) {
            std::complex<double> gammatsm=0.;
            double vargammatsm=0.,crosswtsm=0.;

            for(int j=0;j<int(crossdata.size());++j) {

                double r = exp(crossdata[j].meanlogr);
                double s = r/R;

                if (s>1/smoothscale && s<smoothscale) {
                    crosswtsm += crossdata[j].weight;
                    gammatsm += crossdata[j].weight * crossdata[j].meangammat;
                    vargammatsm += SQR(crossdata[j].weight) * crossdata[j].vargammat;
                }
            }
            if (crosswtsm > 0) {
                gammatsm /= crosswtsm;
                vargammatsm /= SQR(crosswtsm);
            }
            Assert(vargammatsm >= 0.);
            fout << sci(real(gammatsm)) << sci(sqrt(vargammatsm));
            fout << sci(imag(gammatsm)) << sci(sqrt(vargammatsm));
        } else {
            std::complex<double> gammat=crossdata[i].meangammat;
            double vargammat=crossdata[i].vargammat;
            Assert(vargammat >= 0.);
            fout << sci(real(gammat)) << sci(sqrt(vargammat));
            fout << sci(imag(gammat)) << sci(sqrt(vargammat));
        }

        fout << std::endl;
    }
}

void WriteNorm(std::ostream& fout, double minsep, double binsize,
               const std::vector<BinData2<NData,EData> >& crossdata,
               const std::vector<BinData2<EData,EData> >& twoptdata,
               const std::vector<BinData2<NData,NData> >& dd,
               const std::vector<BinData2<NData,NData> >& dr,
               const std::vector<BinData2<NData,NData> >& rr)
{
    Form sci;
    Form fix; 
    sci.sci().prec(3).width(11).right().trail(1); 
    fix.fix().prec(4).width(11).right().trail(1);

    Assert(crossdata.size() ==  twoptdata.size());
    Assert(crossdata.size() ==  dd.size());
    Assert(crossdata.size() ==  dr.size());
    Assert(crossdata.size() ==  rr.size());

    std::vector<double> omega(dd.size());
    std::vector<double> varomega(dd.size());

    for(int i=0;i<int(dd.size());++i) {
        if (rr[i].npair > 0) {
            omega[i] = (dd[i].npair-2*dr[i].npair+rr[i].npair)/rr[i].npair;
            varomega[i] = 1./ rr[i].npair;
        } else {
            omega[i] = varomega[i] = 0.;
        }
    }

    double dlogr = binsize;

    fout <<
        "# R_nominal .    <R>    "
        ".  <NMap>   .   <NMx>   .  sig_nm   "
        ".    <NN>   .   sig_nn  .  <Map^2>  .   sig_mm  "
        ".   nmnorm  . signmnorm .   nnnorm  . signnnorm \n";
    for(int i=0;i<int(crossdata.size());++i) {
        double mapsq=0.,varmap=0.;
        double nmap=0., nmx=0., varnm=0., nn=0., varnn=0.;

        std::complex<double> gammatsm=0.;
        double vargammatsm=0.,crosswtsm=0.;

        double R = exp(twoptdata[i].meanlogr);

        for(int j=0;j<int(crossdata.size());++j) {

            double r = exp(twoptdata[j].meanlogr);
            double xip = real(twoptdata[j].xiplus);
            double xim = real(twoptdata[j].ximinus);
            double var = twoptdata[j].varxi;
            double s = r/R;
            double ssqdlogr = s*s*dlogr;
            double tp = Tplus(s), tm = Tminus(s);

            mapsq += ssqdlogr*(tp*xip + tm*xim)/2.;
            varmap += SQR(ssqdlogr)*(SQR(tp)+SQR(tm))*var/4.;

            double gt = real(crossdata[j].meangammat);
            double gx = imag(crossdata[j].meangammat);
            double tc = Tcross(s);
            nmap += ssqdlogr * tc * gt;
            nmx += ssqdlogr * tc * gx;
            varnm += SQR(ssqdlogr*tc)*crossdata[j].vargammat;

            nn += ssqdlogr * tp * omega[j];
            varnn += SQR(ssqdlogr*tp)*varomega[j];
        }

        if (crosswtsm > 0) {
            gammatsm /= crosswtsm;
            vargammatsm /= SQR(crosswtsm);
        }

        double nmnorm = mapsq*nn == 0. ? 0. : nmap*nmap / (mapsq*nn);
        double varnmnorm = mapsq*nn == 0. ? 0. : SQR(nmap/mapsq/nn)*(4.*varnm) +
            SQR(nmap)*(varnn/SQR(nn) + varmap/SQR(mapsq));

        double nnnorm = mapsq == 0. ? 0. : nn/mapsq;
        double varnnnorm = 
            mapsq == 0. ? 0. :
            varnn/SQR(mapsq) + varmap*SQR(nnnorm/mapsq);

        fout << fix(minsep*exp((i+0.5)*binsize)) << fix(R);
        Assert(varnm >= 0.);
        fout << sci(nmap) << sci(nmx) << sci(sqrt(varnm));
        Assert(varnn >= 0.);
        fout << sci(nn) << sci(sqrt(varnn));
        Assert(varmap >= 0.);
        fout << sci(mapsq) << sci(sqrt(varmap));
        Assert(varnmnorm >= 0.);
        fout << sci(nmnorm) << sci(sqrt(varnmnorm));
        Assert(varnnnorm >= 0.);
        fout << sci(nnnorm) << sci(sqrt(varnnnorm));
        fout << std::endl;
    }
}

void WriteNN(std::ostream& fout, double minsep, double binsize,
             const std::vector<BinData2<NData,NData> >& dd,
             const std::vector<BinData2<NData,NData> >& dr, 
             const std::vector<BinData2<NData,NData> >& rd, 
             const std::vector<BinData2<NData,NData> >& rr)
{
    Form sci;
    Form fix; 
    sci.sci().prec(3).width(11).right().trail(1); 
    fix.fix().prec(4).width(11).right().trail(1);

    Assert(dd.size() == dr.size());
    Assert(dd.size() == rd.size());
    Assert(dd.size() == rr.size());

    std::vector<double> omega(dd.size());
    std::vector<double> varomega(dd.size());

    for(int i=0;i<int(dd.size());++i) {
        if (rr[i].npair > 0) {
            omega[i] = (dd[i].npair-dr[i].npair-rd[i].npair+rr[i].npair)/rr[i].npair;
            varomega[i] = 1./rr[i].npair;
        } else {
            omega[i] = varomega[i] = 0.;
        }
    }

    fout <<
        "# R_nominal .    <R>    "
        ".  omega    . sig_omega .     DD    .    DR     .    RD     .    RR     \n";
    for(int i=0;i<int(omega.size());++i) {

        double R = exp(dd[i].meanlogr);
        fout << fix(minsep*exp((i+0.5)*binsize)) << fix(R);
        Assert(varomega[i] >= 0.);
        fout << sci(omega[i]) << sci(sqrt(varomega[i]));
        fout << sci(dd[i].npair) << sci(dr[i].npair);
        fout << sci(rd[i].npair) << sci(rr[i].npair);
        fout << std::endl;
    }
}


